#!/usr/bin/env python
# -*- coding: utf-8 -*-
#只是识别绿色的无人机测试

from traceback import print_tb
import numpy as np
import rclpy
from rclpy.node import Node
from sensor_msgs.msg import Image
from geometry_msgs.msg import Twist
import argparse
import cv2
import sys
from djitellopy import Tello
import math
# import cv2.aruco

tello = Tello()
class RAS_Tello(Node):
   def __init__(self):
       super().__init__('tello_project')

       self.subscription = self.create_subscription(
           Image,
           '/image_raw',
           self.image_callback,
           10
       )

       self.publisher_ = self.create_publisher(Twist, '/control', 10)
       self.tello = Tello()

   def image_callback(self, msg):

       self.frame = self.imgmsg_to_cv2(msg)   ###-------------------------------------------------
       # RGB to HSV color space
       hsv = cv2.cvtColor(self.frame, cv2.COLOR_BGR2HSV)
       # structural element
       line = cv2.getStructuringElement(cv2.MORPH_RECT, (15, 15), (-1, -1))
       # HSV
       mask_white = cv2.inRange(hsv, (0,0,221), (180,30,255))
       # mask_blue = cv2.inRange(hsv, (100, 43, 46), (124,255,255))
       mask_green = cv2.inRange(hsv, (35,43,46), (77,255,255))
       # mask_yellow = cv2.inRange(hsv,(26,43,46),(34,255,255))
       mask = mask_green
       #Contour extraction, find the largest contour
       i = 1
       # for mask in masks:

       contours, hierarchy = cv2.findContours(mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
       index = -1
       max = 0
       for c in range(len(contours)):
           area = cv2.contourArea(contours[c])
           if area > max:
               max = area
               index = c
       # pic
       if index >= 0:
           rect = cv2.minAreaRect(contours[index])
           #Ellipse Fitting
           cv2.ellipse(self.frame, rect, (0, 255, 0), 2, 8)
           #center point positioning
           cv2.circle(self.frame, (np.int32(rect[0][0]), np.int32(rect[0][1])), 2, (0, 255, 0), 2, 8, 0)
       if i == 1:
           text = 'green'
       # elif i == 2:
       #     text = 'blue'
       # elif i == 3:
       #     text = 'green'
       # elif i == 4:
       #     text = 'yellow'
       # cv2.putText(self.frame,text,(np.int32(rect[0][0]), np.int32(rect[0][1]) ), cv2.FONT_HERSHEY_SIMPLEX, 0.75, (0, 0, 255), 2)
       # i += 1
       cv2.imshow("Frame", self.frame)
       cv2.waitKey(1)


   def tello_movement(self):
       msg = Twist()
       print("[INFO] moving Tello...")
       self.publisher_.publish(msg)
       print('[INFO] Tello moved')

   def imgmsg_to_cv2(self, img_msg):
       n_channels = len(img_msg.data) // (img_msg.height * img_msg.width)
       dtype = np.uint8

       img_buf = np.asarray(img_msg.data, dtype=dtype) if isinstance(img_msg.data, list) else img_msg.data

       if n_channels == 1:
           cv2_img = np.ndarray(shape=(img_msg.height, img_msg.width), dtype=dtype, buffer=img_buf)
       else:
           cv2_img = np.ndarray(shape=(img_msg.height, img_msg.width, n_channels), dtype=dtype, buffer=img_buf)

       # If the byte order is different between the message and the system
       if img_msg.is_bigendian == (sys.byteorder == 'little'):
           cv2_img = cv2_img.byteswap().newbyteorder()

       return cv2_img

def main(args=None):
   rclpy.init(args=args)
   tello_project = RAS_Tello()
   rclpy.spin(tello_project)
   tello_project.destroy_node()
   rclpy.shutdown()

if __name__ == '__main__':
   main()